﻿///////////////////////////////////////////////////////////////////////////////
//
//  This file is part of MathLib.NET.
//
//  This library is free software; you can redistribute it and/or
//  modify it under the terms of the GNU Lesser General Public
//  License as published by the Free Software Foundation; either
//  version 2.1 of the License, or (at your option) any later version.
//  
//  This library is distributed in the hope that it will be useful,
//  but WITHOUT ANY WARRANTY; without even the implied warranty of
//  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
//  Lesser General Public License for more details.
//  
//  You should have received a copy of the GNU Lesser General Public
//  License along with this library;  If not, see 
//  <http://www.gnu.org/licenses/>.
//
///////////////////////////////////////////////////////////////////////////////

using System;
using ILNumerics;
using ILNumerics.Misc;
using MathNetTransforms = MathNet.Numerics.IntegralTransforms.Algorithms;

// Fourier Transforms functions
// ============================
//
// abs - Absolute value and complex magnitude 
// angle - Phase angle 
// cplxpair - Sort complex numbers into complex conjugate pairs 
// fft - Discrete Fourier transform 
// fft2 - 2-D discrete Fourier transform 
// fftn - N-D discrete Fourier transform 
// fftshift - Shift zero-frequency component to center of spectrum 
// fftw - Interface to FFTW library run-time algorithm tuning control 
// ifft - Inverse discrete Fourier transform 
// ifft2 - 2-D inverse discrete Fourier transform 
// ifftn - N-D inverse discrete Fourier transform 
// ifftshift - Inverse FFT shift 
// nextpow2 - Next higher power of 2 
// unwrap - Correct phase angles to produce smoother phase plots

namespace MathLib
{
    public partial class MLMath
    {
        /// <summary>
        /// Object for performing real FFTs.
        /// </summary>
        private static MathNetTransforms::DiscreteFourierTransform mDiscreteFft = new MathNetTransforms.DiscreteFourierTransform();

        /// <summary>
        /// 2-D FFT of real data.
        /// </summary>
        /// <param name="input">Real time-series input</param>
        /// <param name="numXPoints">Number of points in first dimension</param>
        /// <param name="numYPoints">Number of points in second dimension</param>
        /// <returns></returns>
        public static ILArray<complex> fft2(ILArray<double> input, int numXPoints = -1, int numYPoints = -1)
        {
            return fft(fft(input, numXPoints != -1 ? (object)numXPoints : null, 0), numYPoints != - 1 ? (object)numYPoints : null, 1);
        }

        /// <summary>
        /// 2-D FFT of complex data.
        /// </summary>
        /// <param name="input">Complex time-series input</param>
        /// <param name="numXPoints">Number of points in first dimension</param>
        /// <param name="numYPoints">Number of points in second dimension</param>
        /// <returns></returns>
        public static ILArray<complex> fft2(ILArray<complex> input, int numXPoints = -1, int numYPoints = -1)
        {
            return fft(fft(input, numXPoints != -1 ? (object)numXPoints : null, 0), numYPoints != -1 ? (object)numYPoints : null, 1);
        }

        /// <summary>
        /// FFT of real data.
        /// </summary>
        /// <param name="input">Real time series input</param>
        /// <param name="outputPoints">Number of points to transform</param>
        /// <param name="dim">Dimension to transform along</param>
        /// <returns>Complex frequency series output</returns>
        public static ILArray<complex> fft(ILArray<double> input, object outputPoints = null, int dim = -1)
        {
            // Calculate parameters for accessing array
            int numTransforms, transformStride, numPoints, pointStride;
            ArrayHelper.CalcArrayAccessParams<double>(input, ref dim, out numTransforms, out transformStride, out numPoints, out pointStride);

            // Calculate number of output points
            int outPoints = outputPoints is int ? (int)outputPoints : numPoints;
            int transformStrideOut = transformStride / numPoints * outPoints;

            // Determine number of points to copy from input
            int inPoints = numPoints < outPoints ? numPoints : outPoints;

            // Create output array
            int[] outputDims = input.Dimensions.ToIntArray();
            outputDims[dim] = outPoints;
            ILArray<complex> output = new ILArray<complex>(outputDims);

            // Process pairs of transforms
            int minorTransform = 0;
            int majorTransform = 0;
            int numTransformsDiv2 = numTransforms / 2;
            for (int t = 0; t < numTransformsDiv2; t++)
            {
                // Calculate offset of first transform point within flat array.
                int transformOfsIn1 = majorTransform * transformStride + minorTransform;
                int transformOfsOut1 = majorTransform * transformStrideOut + minorTransform;
                minorTransform++;
                if (minorTransform == pointStride)
                {
                    minorTransform = 0;
                    majorTransform++;
                }
                int transformOfsIn2 = majorTransform * transformStride + minorTransform;
                int transformOfsOut2 = majorTransform * transformStrideOut + minorTransform;
                minorTransform++;
                if (minorTransform == pointStride)
                {
                    minorTransform = 0;
                    majorTransform++;
                }

                // Set even transform as real and odd transform as imaginary
                System.Numerics.Complex[] complexInput = new System.Numerics.Complex[outPoints];
                for (int p = 0; p < inPoints; p++)
                {
                    complexInput[p] = new System.Numerics.Complex(
                        input.InternalArray4Experts[transformOfsIn1 + p * pointStride],
                        input.InternalArray4Experts[transformOfsIn2 + p * pointStride]);
                }

                // Perform complex FFT
                mDiscreteFft.BluesteinForward(complexInput, MathNet.Numerics.IntegralTransforms.FourierOptions.Matlab);

                // Extract two real FFTs
                for (int p = 1; p < outPoints; p++)
                {
                    System.Numerics.Complex conjNmk = System.Numerics.Complex.Conjugate(complexInput[outPoints - p]);
                    System.Numerics.Complex fft1 = 0.5 * (complexInput[p] + conjNmk);
                    System.Numerics.Complex fft2 = -0.5 * System.Numerics.Complex.ImaginaryOne * (complexInput[p] - conjNmk);
                    output.InternalArray4Experts[transformOfsOut1 + p * pointStride].real = fft1.Real;
                    output.InternalArray4Experts[transformOfsOut1 + p * pointStride].imag = fft1.Imaginary;
                    output.InternalArray4Experts[transformOfsOut2 + p * pointStride].real = fft2.Real;
                    output.InternalArray4Experts[transformOfsOut2 + p * pointStride].imag = fft2.Imaginary;
                }
                output.InternalArray4Experts[transformOfsOut1].real = complexInput[0].Real;
                output.InternalArray4Experts[transformOfsOut1].imag = 0.0;
                output.InternalArray4Experts[transformOfsOut2].real = complexInput[0].Imaginary;
                output.InternalArray4Experts[transformOfsOut2].imag = 0.0;
            }

            // Process final transform if number of transforms was odd
            if (numTransformsDiv2 * 2 < numTransforms)
            {
                // Calculate offset of first transform point within flat array.
                int transformOfs = majorTransform * transformStride + minorTransform;

                // Set even transform as real and odd transform as imaginary
                System.Numerics.Complex[] complexInput = new System.Numerics.Complex[outPoints];
                for (int p = 0; p < inPoints; p++)
                {
                    complexInput[p] = new System.Numerics.Complex(
                        input.InternalArray4Experts[transformOfs + p * pointStride],
                        0.0);
                }

                // Perform complex FFT
                mDiscreteFft.BluesteinForward(complexInput, MathNet.Numerics.IntegralTransforms.FourierOptions.Matlab);

                // Copy to output
                for (int p = 0; p < outPoints; p++)
                {
                    output.InternalArray4Experts[transformOfs + p * pointStride].real = complexInput[p].Real;
                    output.InternalArray4Experts[transformOfs + p * pointStride].imag = complexInput[p].Imaginary;
                }
            }

            return output;
        }

        /// <summary>
        /// FFT of complex data.
        /// </summary>
        /// <param name="input">Complex time series input</param>
        /// <param name="outputPoints">Number of points to transform</param>
        /// <param name="dim">Dimension to transform along</param>
        /// <returns>Complex frequency series output</returns>
        public static ILArray<complex> fft(ILArray<complex> input, object outputPoints = null, int dim = -1)
        {
            // Calculate parameters for accessing array
            int numTransforms, transformStride, numPoints, pointStride;
            ArrayHelper.CalcArrayAccessParams<complex>(input, ref dim, out numTransforms, out transformStride, out numPoints, out pointStride);

            // Calculate number of output points
            int outPoints = outputPoints is int ? (int)outputPoints : numPoints;
            int transformStrideOut = transformStride / numPoints * outPoints;

            // Determine number of points to copy from input
            int inPoints = numPoints < outPoints ? numPoints : outPoints;

            // Create output array
            int[] outputDims = input.Dimensions.ToIntArray();
            outputDims[dim] = outPoints;
            ILArray<complex> output = new ILArray<complex>(outputDims);

            // Process pairs of transforms
            int minorTransform = 0;
            int majorTransform = 0;
            for (int t = 0; t < numTransforms; t++)
            {
                // Calculate offset of first transform point within flat array.
                int transformOfsIn = majorTransform * transformStride + minorTransform;
                int transformOfsOut = majorTransform * transformStrideOut + minorTransform;
                minorTransform++;
                if (minorTransform == pointStride)
                {
                    minorTransform = 0;
                    majorTransform++;
                }

                // Set even transform as real and odd transform as imaginary
                System.Numerics.Complex[] complexInput = new System.Numerics.Complex[outPoints];
                for (int p = 0; p < inPoints; p++)
                {
                    complexInput[p] = new System.Numerics.Complex(
                        input.InternalArray4Experts[transformOfsIn + p * pointStride].real,
                        input.InternalArray4Experts[transformOfsIn + p * pointStride].imag);
                }

                // Perform complex FFT
                mDiscreteFft.BluesteinForward(complexInput, MathNet.Numerics.IntegralTransforms.FourierOptions.Matlab);

                // Copy to output
                for (int p = 0; p < outPoints; p++)
                {
                    output.InternalArray4Experts[transformOfsOut + p * pointStride].real = complexInput[p].Real;
                    output.InternalArray4Experts[transformOfsOut + p * pointStride].imag = complexInput[p].Imaginary;
                }
            }

            return output;
        }

        /// <summary>
        /// FFT of real data.
        /// </summary>
        /// <param name="input">Real time series input</param>
        /// <param name="outputPoints">Number of points to transform</param>
        /// <param name="dim">Dimension to transform along</param>
        /// <returns>Complex frequency series output</returns>
        public static ILArray<complex> fft(ILArray<float> input, object outputPoints = null, int dim = -1)
        {
            return fft(ILNumerics.BuiltInFunctions.ILMath.todouble(input), outputPoints, dim);
        }

        /// <summary>
        /// FFT of complex data.
        /// </summary>
        /// <param name="input">Complex time series input</param>
        /// <param name="outputPoints">Number of points to transform</param>
        /// <param name="dim">Dimension to transform along</param>
        /// <returns>Complex frequency series output</returns>
        public static ILArray<complex> fft(ILArray<fcomplex> input, object outputPoints = null, int dim = -1)
        {
            return fft(ILNumerics.BuiltInFunctions.ILMath.tocomplex(input), outputPoints, dim);
        }

        /// <summary>
        /// Inverse-FFT of complex data.
        /// </summary>
        /// <param name="input">Complex frequency series input</param>
        /// <returns>Complex time series output</returns>
        public static ILArray<complex> ifft(ILArray<complex> input, object outputPoints = null, int dim = -1)
        {
            // Calculate parameters for accessing array
            int numTransforms, transformStride, numPoints, pointStride;
            ArrayHelper.CalcArrayAccessParams<complex>(input, ref dim, out numTransforms, out transformStride, out numPoints, out pointStride);

            // Calculate number of output points
            int outPoints = outputPoints is int ? (int)outputPoints : numPoints;
            int transformStrideOut = transformStride / numPoints * outPoints;

            // Determine number of points to copy from input
            int inPoints = numPoints < outPoints ? numPoints : outPoints;

            // Create output array
            int[] outputDims = input.Dimensions.ToIntArray();
            outputDims[dim] = outPoints;
            ILArray<complex> output = new ILArray<complex>(outputDims);

            // Process pairs of transforms
            int minorTransform = 0;
            int majorTransform = 0;
            for (int t = 0; t < numTransforms; t++)
            {
                // Calculate offset of first transform point within flat array.
                int transformOfsIn = majorTransform * transformStride + minorTransform;
                int transformOfsOut = majorTransform * transformStrideOut + minorTransform;
                minorTransform++;
                if (minorTransform == pointStride)
                {
                    minorTransform = 0;
                    majorTransform++;
                }

                // Set even transform as real and odd transform as imaginary
                System.Numerics.Complex[] complexInput = new System.Numerics.Complex[outPoints];
                for (int p = 0; p < inPoints; p++)
                {
                    complexInput[p] = new System.Numerics.Complex(
                        input.InternalArray4Experts[transformOfsIn + p * pointStride].real,
                        input.InternalArray4Experts[transformOfsIn + p * pointStride].imag);
                }

                // Perform complex FFT
                mDiscreteFft.BluesteinInverse(complexInput, MathNet.Numerics.IntegralTransforms.FourierOptions.Matlab);

                // Copy to output
                for (int p = 0; p < outPoints; p++)
                {
                    output.InternalArray4Experts[transformOfsOut + p * pointStride].real = complexInput[p].Real;
                    output.InternalArray4Experts[transformOfsOut + p * pointStride].imag = complexInput[p].Imaginary;
                }
            }

            return output;
        }
    }
}
